#include <pthread.h>
#include <mutex>
#include <thread>
#include <unistd.h>
#include <signal.h>
#include "opencv2/opencv.hpp"

#include "cnn_function.hpp"
#include "util.h"
#include "freetype.hpp"

#include "rpc/server.h"

std::vector<std::string> current_user_name;
std::vector<int> current_user_spoofing;
std::vector<cv::Rect> rois_global;

cv::Ptr<cv::freetype::FreeType2> ft2;
const static cv::Scalar default_color = cv::Scalar(0x00, 0x00, 0xff);
bool have_init = false;
std::mutex gpu_lock;
cv::Mat registration_img;
bool have_registration_rpc = false;
std::string registration_image_feature_str;
std::mutex camera_lock;
cv::Mat current_camera_image;
bool have_new_image = false;
bool recognition_stop = false;
std::vector<std::string> user_name;
std::vector<std::vector<float>> user_feature;
bool feature_db_reload = false;
std::string feature_db_path = "/home/nvidia/.bin/git/registration_online/feature_db.txt";

void reload_feature_db(){
    user_name.clear();
    user_feature.clear();
    FILE *fp = fopen (feature_db_path.c_str(), "r");
    if(fp == NULL){
        printf("reload_feature_db: feature_db not exist !\n");
        return;
    }
    char *line = NULL;
    size_t len = 0;
    ssize_t read;
    int index = 1;
    printf("reload_feature_db start!\n");
    while ((read = getline(&line, &len, fp)) != -1) {
        //printf("Retrieved line of length %zu:\n", read);
        char *current_name_str;
        if((current_name_str = strchr(line, ' ')) != NULL) *current_name_str = '\0';
        std::string current_name(line);
        read = getline(&line, &len, fp);
        if(read == -1){
            printf("reload_feature_db over!\n");
            break;
        }
        std::vector<float> current_feature;
        char *line_bak = line;
        for(int i = 0; i < 512; i++){
            current_feature.push_back(atof(line_bak));
            line_bak = strchr(line_bak, ',')+1;
        }
        user_feature.push_back(current_feature);
        user_name.push_back(current_name);
        printf("%d %s\n", index, current_name.c_str());
        index += 1;
    }
    free(line);
    fclose(fp);
}

void free_dnn_network(){
    uninit_detector();
    uninit_recognition();
    uninit_mtcnn();
}

void init_dnn_network(){
    have_init = true;
    reload_feature_db();
    init_detector("/home/nvidia/.bin/git/cnn/cfg/yolov3-small.cfg", "/home/nvidia/.bin/git/cnn/model/yolov3-small_final_max_epoch_15.weights");
    init_recognition("/home/nvidia/.bin/git/cnn/cfg/cosface_new.cfg", "/home/nvidia/.bin/git/cnn/model/model.cnn.50");
    init_mtcnn("/home/nvidia/.bin/git/cnn/cfg/mtcnn_onet.cfg", "/home/nvidia/.bin/git/cnn/model/mtcnn_final.weights");
}

std::string registration_image(std::vector<char> const& img, int reload_db)
{
    if(reload_db == 1){
        std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
        reload_feature_db();
        return std::to_string(0);
    }
    //double start = getMsecOfNow();
    cv::Mat image_input = cv::imdecode(img, CV_LOAD_IMAGE_COLOR);
    {
        std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
        have_registration_rpc = true;
        registration_img = image_input;
    }
    while(have_registration_rpc){
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    //printf("registration_image over spend: %f ms\n", getMsecOfNow() - start);
    return registration_image_feature_str;
}

void get_registration_image_feature(){
    int code = 0;
    cv::Mat image_input = registration_img;

    if(!have_init){
        init_dnn_network();
    }
    int detection_bbox[MAX_BBOX_NUM * 4];

    std::vector<cv::Mat> aligned_faces;
    int face_features_length = 512;
    double start = getMsecOfNow();
    int face_count = 0;
    detection_face(image_input, &face_count, detection_bbox);
    //printf("detection_face spend: %f ms\n", getMsecOfNow() - start);
    printf("registration_image detection_face over face_count %d image width %d, height %d\n",
           face_count, image_input.cols, image_input.rows);
    if(face_count == 0){
        code = 1;  // no face
        registration_image_feature_str = std::to_string(code);
        return;
    } else if(face_count > 1){
        code = 2;  // multi face
        registration_image_feature_str = std::to_string(code);
        return;
    }
    if(detection_bbox[2] - detection_bbox[0] < 200 || detection_bbox[3] - detection_bbox[1] < 200){
        code = 3;  // face size small
        registration_image_feature_str = std::to_string(code);
        return;
    }
    float *face_feature = (float *)malloc(face_count * face_features_length *sizeof(float));
    //for(int i = 0; i < 4; i++) printf("%d %d\n", i, detection_bbox[i]);
    extract_face_feature(image_input, face_count, face_feature, detection_bbox, &code, 1);
    printf("registration_image spend: %f ms code: %d\n", getMsecOfNow() - start, code);
    //for(int i = 0; i < 10; i++) printf("%d %f\n", i, face_feature[i]);
    if(code == 0){
        /*
        FILE *feature_file = fopen (feature_db_path.c_str(), "a");
        fprintf(feature_file, "%s\n", name.c_str());
        for(int i = 0; i < 512; i++){
            fprintf(feature_file, "%f,", face_feature[i]);
        }
        fprintf(feature_file, "\n");
        fclose(feature_file);
        reload_feature_db();
        */
        std::string feature_str = "";
        for(int i = 0; i < 512; i++){
            feature_str += std::to_string(face_feature[i]) + ",";
        }
        registration_image_feature_str = feature_str;
    } else {
        registration_image_feature_str = std::to_string(code);
    }
    free(face_feature);
}

void *registration_thread(void *args) {
    //registration_mysql(1006, 1001);
    rpc::server srv("127.0.0.1", 18080);
    // note: the signature is automatically captured
    /*
    srv.bind("registration_image", [](std::vector<char> const& img, std::string &name) {
            return registration_image(img, name);
        });
    */
    srv.bind("registration_image", &registration_image);
    //srv.run();
    srv.async_run();
    while (!recognition_stop){
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }
    printf("registration_thread quit...\n");
    return 0;
}

float *get_image_feature(cv::Mat &image_input, int *face_count, int *code, int *detection_bbox)
{
    *code = 0;
    if(!have_init){
        init_dnn_network();
    }

    std::vector<cv::Mat> aligned_faces;
    int face_features_length = 512;
    //double start = getMsecOfNow();
    detection_face(image_input, face_count, detection_bbox);
    //printf("detection_face spend: %f ms\n", getMsecOfNow() - start);
    //printf("detection_face over face_count %d\n", *face_count);
    if(*face_count == 0){
        *code = 1;  // no face
        return 0;
    }
    float *face_feature = (float *)malloc(*face_count * face_features_length *sizeof(float));
    //for(int i = 0; i < 4; i++) printf("%d %d\n", i, detection_bbox[i]);

    extract_face_feature(image_input, *face_count, face_feature, detection_bbox, code, 0);
    //printf("run_recognition spend: %f ms code: %d, face_count %d\n", getMsecOfNow() - start, *code, *face_count);
    //for(int i = 0; i < 10; i++) printf("%d %f\n", i, face_feature[i]);
    if(*code == 0){
        return face_feature;
    } else {
        free(face_feature);
        return 0;
    }
}

cv::VideoCapture capture;

void sig_close(int sig){
    const char *s;
    switch (sig) {
    case SIGINT:
        s = "interrupt";
        break;
    case SIGTERM:
        s = "terminate";
        break;
    }
    recognition_stop = true;
    printf("\n\nreceived signal: %s, quit...\n", s);
}

void *camera_thread(void *args){
    capture.open(0);
    if (!capture.isOpened()) {
        printf("Video capture not successfully initialized, exit!\n");
        //return 0;
    }
    capture.set(CV_CAP_PROP_FRAME_WIDTH, 1080);
    capture.set(CV_CAP_PROP_FRAME_HEIGHT, 960);
    capture.set(CV_CAP_PROP_FPS, 25);
    cv::Mat img;
    int not_captured_image = 0;
    while(!recognition_stop){
        bool bCap = true;
        bCap = capture.read(img);
        if (!bCap) {
            printf("warning: frame not captured: %d\n", not_captured_image);
            not_captured_image += 1;
            if(not_captured_image > 10){
                recognition_stop = true;
                break;
            }
            usleep(20000);
            continue;
        }
        not_captured_image = 0;
        {
            std::lock_guard<std::mutex> camera_lock_guard(camera_lock, std::adopt_lock);
            current_camera_image = img.clone();
            have_new_image = true;
        }
        usleep(20000);
    }
    capture.release();
    printf("camera_thread quit.\n");
    return 0;
}

int main(int argc, char **argv)
{
    signal(SIGINT, sig_close);  //close CTRL+C
    signal(SIGTERM, sig_close);
    ft2 = cv::freetype::createFreeType2();
    ft2->loadFontData("/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc", 0);
    
    cv::Mat image_input;
    cv::namedWindow("Face Recognition", CV_WINDOW_NORMAL);
    cv::setWindowProperty("Face Recognition", CV_WND_PROP_FULLSCREEN, CV_WINDOW_FULLSCREEN);

    pthread_t recognition_thread_pid;
    pthread_create(&recognition_thread_pid, NULL, registration_thread, NULL);
    pthread_t camera_thread_pid;
    pthread_create(&camera_thread_pid, NULL, camera_thread, NULL);
    int detection_bbox[MAX_BBOX_NUM * 4];
    while(!recognition_stop){
        if(have_registration_rpc){
            get_registration_image_feature();
            have_registration_rpc = false;
        }
        while(!recognition_stop){
            {
                std::lock_guard<std::mutex> camera_lock_guard(camera_lock, std::adopt_lock);
                if(have_new_image){
                    image_input = current_camera_image.clone();
                    have_new_image = false;
                    break;
                }
            }
            usleep(10000);
        }
        int face_count = 0;
        int code = -1;
        float *face_feature;
        {
            std::lock_guard<std::mutex> gpu_lock_guard(gpu_lock, std::adopt_lock);
            face_feature = get_image_feature(image_input, &face_count, &code, detection_bbox);
        }
        if(code == 0){
            for(int m = 0; m < face_count; m++){
                cv::Rect roi_local(cv::Point(detection_bbox[m * 4], detection_bbox[m * 4 + 1]),
                                   cv::Size(detection_bbox[m * 4 + 2] - detection_bbox[m * 4],
                                            detection_bbox[m * 4 + 3] - detection_bbox[m * 4 + 1]));
                cv::rectangle(image_input, cvPoint(roi_local.x, roi_local.y),
                              cvPoint(roi_local.x + roi_local.width, roi_local.y + roi_local.height), cvScalar(255,0,0), 4);
                int max_score_index = -1;
                float max_score = -1.0f;
                for(unsigned int i = 0; i < user_feature.size(); i++){
                    float current_score = 0;
                    for(int j = 0; j < 512; j++){
                        current_score += face_feature[m * 512 + j] * user_feature[i][j];
                    }
                    if(current_score > max_score){
                        max_score = current_score;
                        max_score_index = i;
                    }
                }
                if(max_score > 0.40f){
                    printf("recognition success: %s, score %f\n", user_name[max_score_index].c_str(), max_score);
                    cv::Point textOrg(roi_local.x, roi_local.y + image_input.rows / 15);
                    int thickness = -1;
                    int linestyle = 8;
                    ft2->putText(image_input, user_name[max_score_index], textOrg, image_input.rows / 15,
                                 default_color, thickness, linestyle, true);
                }
            }
            free(face_feature);
        }
        //cv::imwrite("/home/nvidia/image.jpg", image_input);
        cv::imshow("Face Recognition", image_input);
        cv::waitKey(5);
        //usleep(100000);
    }
    printf("main quit...\n");
    pthread_join(recognition_thread_pid, NULL);
    pthread_join(camera_thread_pid, NULL);
    free_dnn_network();
    return 0;
}
